#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2018/12/19 21:04
# @USER    : Connor
# @File    : CineSubmission_MAC.py
# @Software: PyCharm
# @Version  : Python-3.6
# @TASK:
import logging
import cv2
import scipy.io as sio
import matplotlib.pyplot as plt
import os,glob,re
import numpy as np
from helpers import reshape,center_crop

def T1mapping_display():
    path = os.path.join(os.getcwd(),'T1MAPPING')
    filename = os.path.join(path,'shaozhongzhou2.npy')
    endomasks = np.load(os.path.join(path,'i_pred_masks.npy'))
    epimasks = np.load(os.path.join(path,'o_pred_masks.npy'))
    images = np.load(filename)
    h, w, depth = images.shape
    crop_size = 100
    for time in range(7):
        for slice in range(11):
            idx = time*11+slice
            img = images[:, :, idx]
            if img.ndim<3:
                img = img[...,np.newaxis]
            itmp = reshape(endomasks[idx], to_shape=(h, w, 1))
            otmp = reshape(epimasks[idx], to_shape=(h, w, 1))
            itmp = np.where(itmp > 0.5, 255, 0).astype('uint8')
            otmp = np.where(otmp > 0.5, 255, 0).astype('uint8')
            assert img.shape == itmp.shape, 'Prediction does not match shape'
            assert img.shape == otmp.shape, 'Prediction does not match shape'
            tmp2, coords, hierarchy = cv2.findContours(otmp.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
            if len(coords) > 1:
                lengths = []
                for coord in coords:
                    lengths.append(len(coord))
                coords = coords[np.argmax(lengths)]
            # for ocoord in coords:
            ocoord = np.squeeze(coords, axis=(1,))
            ocoord = np.append(ocoord, ocoord[:1], axis=0)
            tmp2, coords, hierarchy = cv2.findContours(itmp.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
            if len(coords) > 1:
                lengths = []
                for coord in coords:
                    lengths.append(len(coord))
                coords = coords[np.argmax(lengths)]
            # for icoord in coords:
            icoord = np.squeeze(coords, axis=(1,))
            icoord = np.append(icoord, icoord[:1], axis=0)


            plt.figure(figsize=(10, 5))
            plt.axis("off")
            plt.imshow(img[:, :,0], cmap=plt.cm.gray)
            plt.plot(icoord[:, 0], icoord[:, 1],color='r')
            plt.plot(ocoord[:, 0], ocoord[:, 1],color='g')
            if time<3:
                filename = 'Native_{:d}_Time_{:d}'.format(time+1,slice+1)
            else:
                filename = 'Post_{:d}_Time_{:d}'.format(time+1,slice+1)
            plt.title(filename)
            plt.savefig(os.path.join(path,filename), bbox_inches='tight')
            plt.close()
            # plt.show()

    # for idx in range(depth):
    #     img = images[:, :, idx]
    #     if img.ndim<3:
    #         img = img[...,np.newaxis]
    #     itmp = reshape(endomasks[idx], to_shape=(h, w, 1))
    #     otmp = reshape(epimasks[idx], to_shape=(h, w, 1))
    #     itmp = np.where(itmp > 0.5, 255, 0).astype('uint8')
    #     otmp = np.where(otmp > 0.5, 255, 0).astype('uint8')
    #     assert img.shape == itmp.shape, 'Prediction does not match shape'
    #     assert img.shape == otmp.shape, 'Prediction does not match shape'
    #     tmp2, coords, hierarchy = cv2.findContours(otmp.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
    #     if len(coords) > 1:
    #         lengths = []
    #         for coord in coords:
    #             lengths.append(len(coord))
    #         coords = coords[np.argmax(lengths)]
    #     # for ocoord in coords:
    #     ocoord = np.squeeze(coords, axis=(1,))
    #     ocoord = np.append(ocoord, ocoord[:1], axis=0)
    #     tmp2, coords, hierarchy = cv2.findContours(itmp.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
    #     if len(coords) > 1:
    #         lengths = []
    #         for coord in coords:
    #             lengths.append(len(coord))
    #         coords = coords[np.argmax(lengths)]
    #     # for icoord in coords:
    #     icoord = np.squeeze(coords, axis=(1,))
    #     icoord = np.append(icoord, icoord[:1], axis=0)
    #
    #
    #     plt.figure(figsize=(10, 5))
    #     plt.axis("off")
    #     plt.imshow(img[:, :,0], cmap=plt.cm.gray)
    #     plt.plot(icoord[:, 0], icoord[:, 1],color='r')
    #     plt.plot(ocoord[:, 0], ocoord[:, 1],color='g')
    #     plt.show()

def display():
    directory = os.path.join(os.getcwd(),'MAC-Cine')
    glob_search = os.path.join(directory,'3P00*')
    subdirs = glob.glob(glob_search)
    if len(subdirs)==0:
        logger.error("Couldn't find Group listing file in {}. "
                            "Wrong directory?".format(directory))
        raise Exception
    subdirs.reverse()
    for idx,subdir in enumerate(subdirs):
        glob_search = os.path.join(subdir, '*_sBTFE_BH_7_*.mat')
        matfiles = glob.glob(glob_search)
        for file in matfiles:
            data = sio.loadmat(file)
            images = data['image']
            # timages = timages[0, 0]
            images = np.array(images)
        # if idx==0:
        #     images=timages
        # else:
        #     images = np.concatenate((images,timages),axis=2)
    h, w, depth = images.shape
    plt.figure(figsize=(10, 5))
    for index in range(depth):
        plt.axis("off")
        plt.imshow(images[:,:,index], cmap=plt.cm.gray)
    # contourdir = os.path.join(directory,'FCN_auto_contour20190109')
    # endomasks=np.load(os.path.join(contourdir,'i_pred_masks.npy'))
    # epimasks=np.load(os.path.join(contourdir,'o_pred_masks.npy'))
    # h,w,depth = images.shape
    #
    # for idx in range(depth):
    #     img = images[:,:,idx]
    #     if img.ndim < 3:
    #         img = img[..., np.newaxis]
    #     itmp = reshape(endomasks[idx], to_shape=(h, w, 1))
    #     otmp = reshape(epimasks[idx], to_shape=(h, w, 1))
    #     itmp = np.where(itmp > 0.5, 255, 0).astype('uint8')
    #     otmp = np.where(otmp > 0.5, 255, 0).astype('uint8')
    #     assert img.shape == itmp.shape, 'Prediction does not match shape'
    #     assert img.shape == otmp.shape, 'Prediction does not match shape'
    #     tmp2, coords, hierarchy = cv2.findContours(otmp.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
    #     if len(coords) > 1:
    #         lengths = []
    #         for coord in coords:
    #             lengths.append(len(coord))
    #         coords = coords[np.argmax(lengths)]
    #     # for ocoord in coords:
    #     ocoord = np.squeeze(coords, axis=(1,))
    #     ocoord = np.append(ocoord, ocoord[:1], axis=0)
    #     tmp2, coords, hierarchy = cv2.findContours(itmp.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
    #     if len(coords) > 1:
    #         lengths = []
    #         for coord in coords:
    #             lengths.append(len(coord))
    #         coords = coords[np.argmax(lengths)]
    #     # for icoord in coords:
    #     icoord = np.squeeze(coords, axis=(1,))
    #     icoord = np.append(icoord, icoord[:1], axis=0)
    #     # display
    #     plt.figure(figsize=(10, 5))
    #     plt.axis("off")
    #
    #     plt.imshow(img[:,:,0], cmap=plt.cm.gray)
    #     plt.plot(icoord[:, 0], icoord[:, 1],color='r')
    #     plt.plot(ocoord[:, 0], ocoord[:, 1],color='g')
    #     plt.show()
def displayCrop():
    directory = os.path.join(os.getcwd(),'MAC-Cine')
    glob_search = os.path.join(directory,'*.mat')
    subdirs = glob.glob(glob_search)
    if len(subdirs)==0:
        logger.error("Couldn't find Group listing file in {}. "
                            "Wrong directory?".format(directory))
        raise Exception
    # subdirs.reverse()
    for idx,file in enumerate(subdirs):
            data = sio.loadmat(file)
            images = data['image']
            # timages = timages[0, 0]
            images = np.array(images)
        # if idx==0:
        #     images=timages
        # else:
        #     images = np.concatenate((images,timages),axis=2)
    h, w, depth = images.shape
    endomasks=np.load(os.path.join(directory,'i_pred_masks.npy'))
    epimasks=np.load(os.path.join(directory,'o_pred_masks.npy'))
    # plt.figure(figsize=(10, 5))
    # for index in range(depth):
    #     plt.axis("off")
    #     plt.imshow(images[:,:,index], cmap=plt.cm.gray)

    for idx in range(depth):
        img = images[:,:,idx]
        if img.ndim < 3:
            img = img[..., np.newaxis]
        itmp = reshape(endomasks[idx], to_shape=(h, w, 1))
        otmp = reshape(epimasks[idx], to_shape=(h, w, 1))
        itmp = np.where(itmp > 0.5, 255, 0).astype('uint8')
        otmp = np.where(otmp > 0.5, 255, 0).astype('uint8')
        assert img.shape == itmp.shape, 'Prediction does not match shape'
        assert img.shape == otmp.shape, 'Prediction does not match shape'
        tmp2, coords, hierarchy = cv2.findContours(otmp.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
        if len(coords) > 1:
            lengths = []
            for coord in coords:
                lengths.append(len(coord))
            coords = coords[np.argmax(lengths)]
        # for ocoord in coords:
        ocoord = np.squeeze(coords, axis=(1,))
        ocoord = np.append(ocoord, ocoord[:1], axis=0)
        tmp2, coords, hierarchy = cv2.findContours(itmp.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
        if len(coords) > 1:
            lengths = []
            for coord in coords:
                lengths.append(len(coord))
            coords = coords[np.argmax(lengths)]
        # for icoord in coords:
        icoord = np.squeeze(coords, axis=(1,))
        icoord = np.append(icoord, icoord[:1], axis=0)
        # display
        plt.figure(figsize=(10, 5))
        plt.axis("off")

        plt.imshow(img[:,:,0], cmap=plt.cm.gray)
        plt.plot(icoord[:, 0], icoord[:, 1],color='r')
        plt.plot(ocoord[:, 0], ocoord[:, 1],color='g')
        plt.show()

def load_CineImages(patientdir):
    glob_search = os.path.join(patientdir, '*_sBTFE_BH_7_*.mat')
    matfiles = glob.glob(glob_search)
    if len(matfiles) == 0:
        logger.error("Couldn't find Cine-Short MAT file in {}. "
                     "Wrong directory?".format(patientdir))
        raise Exception
    for file in matfiles:
        data = sio.loadmat(file)
        images = data['I']
        images = images[0, 0]
        images = np.array(images[0])

    return images

def main():
    directory='E:\\717-3Lab\\MRI\\MAC'
    glob_search = os.path.join(directory,'Group*')
    subdir = glob.glob(glob_search)
    if len(subdir)==0:
        logger.error("Couldn't find Group listing file in {}. "
                            "Wrong directory?".format(directory))
        return
    # import operator
    # # 条件为真，返回前面内容
    # func = lambda x: [y for l in x for y in func(l)] if type(x) is list else [x]
    # patients0 =func([[os.path.join(pdir,p) for p in os.listdir(pdir)] for pdir in subdir])
    patients = [os.path.join(pdir, p) for pdir in subdir for p in os.listdir(pdir)]
    # operator.eq(patients,patients0)
    logger.info('Loading patient images ... ')
    for patient in patients:
        # pass
        images=load_CineImages(patient)

        # # display
        # plt.figure(figsize=(10, 5))
        # for index in range(3):
        #     plt.subplot(1, 3, index+1)
        #     plt.axis("off")
        #     plt.imshow(images[:,:,index], cmap=plt.cm.gray)
        # plt.show()


    logger.info('done!')


if __name__ == '__main__':
    logger = logging.getLogger()
    logger.setLevel(logging.NOTSET)  # Log等级总开关
    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)
    formatter = logging.Formatter("%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s")
    ch.setFormatter(formatter)
    logger.addHandler(ch)
    logger.info('Building Cine Data Set ...')
    # main()
    # display()
    # display1()
    displayCrop()
    # T1mapping_display()